import os
from fastapi import FastAPI
from backend.common.logging import logger
from tortoise import Tortoise
from backend.config.config import settings
from contextlib import asynccontextmanager

# Function to get configuration with environment variables prioritized
def get_db_config():
    return {
        'connections': {
            'default': {
                'engine': f'tortoise.backends.{os.getenv("MYSQL_ENGINE", settings.MYSQL_ENGINE)}',
                'credentials': {
                    'host': os.getenv("MYSQL_HOST", settings.MYSQL_HOST),
                    'port': int(os.getenv("MYSQL_PORT", settings.MYSQL_PORT)),
                    'user': os.getenv("MYSQL_USER", settings.MYSQL_USER),
                    'password': os.getenv("MYSQL_PASSWORD", settings.MYSQL_PASSWORD),
                    'database': os.getenv("MYSQL_DATABASE", settings.MYSQL_DATABASE),
                    'charset': os.getenv("MYSQL_ENCODING", settings.MYSQL_ENCODING),
                    'echo': os.getenv("MYSQL_ECHO", settings.MYSQL_ECHO)
                }
            },
        },
        'apps': {
            'models': {
                'models': [
                    "backend.models.user",
                    "backend.models.request_log",
                    "backend.models.agent"
                ],
                'default_connection': 'default',
            },
        },
        'use_tz': False,
        'timezone': 'Asia/Shanghai'
    }

# Get the database configuration
db_config = get_db_config()

# Add TORTOISE_ORM for aerich migrations
TORTOISE_ORM = get_db_config()

@asynccontextmanager
async def db_lifespan(app: FastAPI):
    await Tortoise.init(config=db_config)
    if settings.GENERATE_SCHEMAS:
        await Tortoise.generate_schemas()
    try:
        conn = Tortoise.get_connection("default")
        await conn.execute_query("SELECT 1")
        logger.info(f"🐬  MySQL connected Successfully")
    except Exception as e:
        logger.error(f"❌ MySQL connected fail: {e}")
        raise
    yield
    await Tortoise.close_connections()
    logger.info("🧹 MySQL closed")
